package com.alibaba.datax.plugin.writer.miniowriter.handler;

import com.alibaba.datax.common.element.Record;
import com.alibaba.datax.common.plugin.RecordReceiver;
import com.alibaba.datax.common.plugin.TaskPluginCollector;
import com.alibaba.datax.common.util.Configuration;
import com.alibaba.datax.plugin.writer.miniowriter.utils.ExcelUtils;
import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.support.ExcelTypeEnum;
import com.alibaba.excel.write.handler.RowWriteHandler;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.alibaba.excel.write.metadata.holder.WriteSheetHolder;
import com.alibaba.excel.write.metadata.holder.WriteTableHolder;
import org.apache.poi.ss.usermodel.Row;

import java.io.InputStream;
import java.io.OutputStream;
import java.util.*;

/**
 * @Author JDragon
 * @Date 2021.07.12 上午 11:36
 * @Email 1061917196@qq.com
 * @Des:
 */
public class ExcelWriterHandler extends WriterHandler {

    public String fileName;

    private OutputStream outputStream;

    public List<List<String>> header;

    private WriteSheet writeSheet;

    private ExcelWriter excelWriter;

    private List<Record> cacheRecord = new LinkedList<>();

    public ExcelWriterHandler(String fileName, OutputStream outputStream) {
        this.fileName = fileName;
        this.outputStream = outputStream;
    }

    public void writerHeader(List<String> header) {
        this.header = new ArrayList<>();
        for (String head : header) {
            List<String> headItem = new ArrayList<>();
            headItem.add(head);
            this.header.add(headItem);
        }
        excelWriter = EasyExcel.write(outputStream).excelType(ExcelTypeEnum.XLSX).build();
        writeSheet = EasyExcel.writerSheet(0, fileName)
                .head(this.header)
                .registerWriteHandler(ExcelUtils.getSheetStyleStrategy())
                .build();
    }

    @Override
    public void write(Record record) {
        cacheRecord.add(record);
        if (cacheRecord.size() >= 1024) {
            writeAll();
        }
    }

    @Override
    public void finish() {
        if (!cacheRecord.isEmpty()) {
            writeAll();
        }
        excelWriter.finish();
    }


    public void writeAll() {
        for (Record r : cacheRecord) {
            try {
                excelWriter.write(Collections.singletonList(resolveRow(r)), writeSheet);
            } catch (Exception e) {
                collectDirtyRecord(r, e);
            }
        }
        cacheRecord.clear();
    }
}
